package encrypt

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"fmt"
)

// AESEncryptECB AES的ECB模式加密
func AESEncryptECB(src []byte, key []byte) []byte {
	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println("new aes cipher err：", err.Error())
	}

	src = PKCS5Padding(src, block.BlockSize()) // padding

	buf := make([]byte, block.BlockSize())
	encrypt := make([]byte, 0)
	for i := 0; i < len(src); i += block.BlockSize() {
		block.Encrypt(buf, src[i:i+block.BlockSize()])
		encrypt = append(encrypt, buf...)
	}

	return encrypt
}

// AESDecryptECB AES的ECB模式解密
func AESDecryptECB(src []byte, key []byte) []byte {

	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println("new aes cipher err：", err.Error())
	}

	buf := make([]byte, block.BlockSize())
	decrypt := make([]byte, 0)
	for i := 0; i < len(src); i += block.BlockSize() {
		block.Decrypt(buf, src[i:i+block.BlockSize()])
		decrypt = append(decrypt, buf...)
	}

	decrypt = PKCS5UnPadding(decrypt)
	return decrypt
}

// AESEncryptCBC AES的CBC模式加密
func AESEncryptCBC(src []byte, key []byte) []byte {
	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println("new aes cipher err：", err.Error())
	}

	src = PKCS5Padding(src, block.BlockSize())
	blockMode := cipher.NewCBCEncrypter(block, key[:block.BlockSize()])
	encrypt := make([]byte, len(src))
	blockMode.CryptBlocks(encrypt, src)
	return encrypt
}

// AESDecryptCBC AES的CBC模式解密
func AESDecryptCBC(src []byte, key []byte) []byte {
	block, err := aes.NewCipher(key)
	if err != nil {
		fmt.Println("new aes cipher err：", err.Error())
	}

	blockMode := cipher.NewCBCDecrypter(block, key[:block.BlockSize()])
	decrypt := make([]byte, len(src))
	blockMode.CryptBlocks(decrypt, src)
	decrypt = PKCS5UnPadding(decrypt)
	return decrypt
}

// PKCS5Padding PKCS5形式的padding
func PKCS5Padding(src []byte, blockSize int) []byte {
	paddingNum := blockSize - len(src)%blockSize
	padding := bytes.Repeat([]byte{byte(paddingNum)}, paddingNum)
	return append(src, padding...)
}

// PKCS5UnPadding PKCS5形式的Unpadding
func PKCS5UnPadding(src []byte) []byte {
	length := len(src)
	paddingNum := int(src[length-1])
	return src[:length-paddingNum]
}

func TestAes() {

	aesKey := []byte("bsslovewcj124723")
	var text = "hello world!"
	ECBcipher := AESEncryptECB([]byte(text), []byte(aesKey))
	ECBplainText := AESDecryptECB(ECBcipher, []byte(aesKey))
	if string(ECBplainText) == text {
		fmt.Println("AES 的ECB加解密模式正确")
	}

	CBCcipher := AESEncryptCBC([]byte(text), []byte(aesKey))
	CBCplainText := AESDecryptCBC(CBCcipher, []byte(aesKey))
	if string(CBCplainText) == text {
		fmt.Println("AES 的CBC加解密模式正确")
	}
}
